#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from utils.detector import train_det, test_det
from utils.mdistance import extr_features, calc_mdistance
from detectron2.engine import default_argument_parser


def parse_opt():
    parser = default_argument_parser()
    
    # ------------------------------------------------------------- #
    # PTL dataset options
    # ------------------------------------------------------------- #
    parser.add_argument('--exp_id', type=str, default='exp1')
    parser.add_argument('--cur_ptl_iter', type=str, default='iter_0')
    parser.add_argument('--syn_smp_angles', type=int, default=30)
    parser.add_argument('--min_bbox_hw_gan', type=int, default=10)
    parser.add_argument('--target_hw_gan', type=int, default=256)

    # ------------------------------------------------------------- #
    # PTL general options
    # ------------------------------------------------------------- #    
    parser.add_argument('--n_gpus', type=int, default=4)
    parser.add_argument('--n_selected_syn_per_iter', type=int, default=100)
    parser.add_argument('--suffix', type=str, default='')
    parser.add_argument('--det_only', action='store_true')
    
    # ------------------------------------------------------------- #
    # PTL detector options
    # ------------------------------------------------------------- #
    parser.add_argument('--batch_size_det', type=int, default=16)
    parser.add_argument('--n_iter_det', type=int, default=6000)
    parser.add_argument('--step_det', type=int, default=5000)
    parser.add_argument('--base_lr_det', type=float, default=0.001)
    parser.add_argument('--pretrained_weights_det', type=str, default='')
    parser.add_argument('--freeze_backbone_det', action='store_true')
    parser.add_argument('--filter_empty_annotations_det', action='store_true')
    parser.add_argument('--aspect_ratio_grouping_det', action='store_true')
    parser.add_argument('--min_size_train_det', nargs='+', type=int, default=[192, 224, 256, 288])
    parser.add_argument('--max_size_train_det', type=int, default=480)
    parser.add_argument('--min_size_train_uav_det', nargs='+', type=int, default=[768, 800, 832, 864])
    parser.add_argument('--max_size_train_uav_det', type=int, default=1440)
    parser.add_argument('--min_size_train_syn_det', nargs='+', type=int, default=[128, 256, 384, 512])
    parser.add_argument('--max_size_train_syn_det', type=int, default=512)
    parser.add_argument('--min_size_test_det', type=int, default=864)
    parser.add_argument('--max_size_test_det', type=int, default=1440)
    parser.add_argument('--progressive', action='store_true')
    parser.add_argument('--checkpoint_period', type=int, default=10000)
    
    # ------------------------------------------------------------- #
    # PTL mdistance options
    # ------------------------------------------------------------- #
    parser.add_argument('--extr_mode', type=str, choices=['single', 'conv'], default='single')
    parser.add_argument('--score_thresh_test_extr', type=float, default=0.1)
    parser.add_argument('--batch_size_extr', type=int, default=32)
    parser.add_argument('--imp_samp', type=str, choices=['direct', 'exp', 'low', 'mid', 'high'], default='exp')
    parser.add_argument('--exp_tau', type=float, default=5.0)
    parser.add_argument('--use_euclidean', action='store_true')
    parser.add_argument('--features_dim', type=int, default=32)
    parser.add_argument('--score_thresh_r', type=float, default=0.3)
    parser.add_argument('--score_thresh_v', type=float, default=0.3)
    parser.add_argument('--md_clip', type=int, default=100)
    parser.add_argument('--display_cdf', action='store_false')
    
    return parser.parse_args()   


def main(opt):
    # Detector training
    train_det(opt)
    
    # Detector testing (intra + inter domains)
    # Remove or add test datasets here
    test_datasets = {'vis': '../datasets/visdrone',
                     'oku': '../datasets/okutama',
                     'icg': '../datasets/icg'}
    for test_dataset_type in test_datasets:
        test_det(opt, test_dataset_type, test_datasets[test_dataset_type])
   
    if not opt.det_only:
        # Feature extraction
        extr_features(opt)
        
        # Mahalanobis distance
        calc_mdistance(opt)


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)